
#include "non_max_suppression_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

constexpr uint32_t BLOCK_SIZE = 32;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    NonMaxSuppressionTilingData tiling;
    const gert::StorageShape* score_shape = context->GetInputShape(1);
    int32_t num_batches = score_shape->GetStorageShape().GetDim(0);
    int32_t num_classes = score_shape->GetStorageShape().GetDim(1);
    int32_t spatial_dimension = score_shape->GetStorageShape().GetDim(2);
    // for (int i = 0; i < score_shape->GetStorageShape().GetDimNum(); i++)
    //     data_sz *= score_shape->GetStorageShape().GetDim(i);
    tiling.set_num_batches(num_batches);
    tiling.set_num_classes(num_classes);
    tiling.set_spatial_dimension(spatial_dimension);

    int32_t max_output_boxes_per_class = 0;
    const gert::Tensor *input_tensor2 = context->GetInputTensor(2);
    if (input_tensor2 == nullptr) {
        return ge::GRAPH_FAILED;
    }
    auto inputaddr2 = input_tensor2->GetData<int32_t>();
    max_output_boxes_per_class = *inputaddr2;
    tiling.set_max_output_boxes_per_class(max_output_boxes_per_class);

    float iou_threshold = 0;
    const gert::Tensor *input_tensor3 = context->GetInputTensor(3);
    if (input_tensor3 == nullptr) {
        return ge::GRAPH_FAILED;
    }
    auto inputaddr3 = input_tensor3->GetData<float>();
    iou_threshold = *inputaddr3;
    tiling.set_iou_threshold(iou_threshold);

    float score_threshold = 0;
    const gert::Tensor *input_tensor4 = context->GetInputTensor(4);
    if (input_tensor4 == nullptr) {
        return ge::GRAPH_FAILED;
    }
    auto inputaddr4 = input_tensor4->GetData<float>();
    score_threshold = *inputaddr4;
    tiling.set_score_threshold(score_threshold);

    auto center_point_boxptr = context->GetAttrs()->GetInt(0);
    tiling.set_center_point_box(*center_point_boxptr);

    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize(); 
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); 
    uint32_t inputLength = inputBytes * inputNum;
    uint32_t ubDataNumber = 12;
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber; //每个ub段可用的空间块数
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes; //每次处理的数据量
    tileDataNum = tileDataNum / BLOCK_SIZE * BLOCK_SIZE;
    tiling.set_tileDataNum(tileDataNum);

    std::cout << "center_point_boxptr:" << *center_point_boxptr << std::endl;
    std::cout << "num_batches:" << num_batches << std::endl;
    std::cout << "num_classes:" << num_classes << std::endl;
    std::cout << "spatial_dimension:" << spatial_dimension << std::endl;
    std::cout << "max_output_boxes_per_class:" << max_output_boxes_per_class << std::endl;
    std::cout << "iou_threshold:" << iou_threshold << std::endl;
    std::cout << "score_threshold:" << score_threshold << std::endl;
    std::cout << "tileDataNum:" << tileDataNum << std::endl;

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* score_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *score_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class NonMaxSuppression : public OpDef {
public:
    explicit NonMaxSuppression(const char* name) : OpDef(name)
    {
        this->Input("boxes")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("scores")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("max_output_boxes_per_class")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("iou_threshold")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("score_threshold")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("selected_indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Attr("center_point_box").AttrType(OPTIONAL).Int(0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(NonMaxSuppression);
}
